import sys
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.applications.densenet import DenseNet121, DenseNet169, DenseNet201

from tensorflow.python.keras.layers import Dense, Flatten, BatchNormalization, MaxPooling2D, GlobalMaxPool2D

from model.activation import Activation
from util import Conf


class Densenet(Conf):
    def __init__(self):
        super(Densenet, self).__init__()
        self.model = self.model_init()

    def model_select(self):
        print(self.model_name)
        if self.model_name == 'DenseNet121':
            return DenseNet121(input_shape=(self.img_width, self.img_height, 3), include_top=False)
        elif self.model_name == 'DenseNet169':
            return DenseNet169(input_shape=(self.img_width, self.img_height, 3), include_top=False)
        elif self.model_name == 'DenseNet201':
            return DenseNet201(input_shape=(self.img_width, self.img_height, 3), include_top=False)
        else:
            print('模型不存在，请确认')
            sys.exit(1)

    def model_init(self):
        model = self.model_select()
        model.trainable = False if self.trainable == 'False' else True
        activation = Activation().get_activation()
        model = Sequential([
            model,
            GlobalMaxPool2D(),
            Flatten(),
            Dense(1000, activation=activation),
            BatchNormalization(),
            Dense(200, activation=activation),
            BatchNormalization(),
            Dense(self.class_number, activation='softmax')])
        return model
